﻿' IndexBuilder.vb
Imports System.Collections.Concurrent
Imports System.IO
Imports System.Text.Json
Imports System.Threading
Imports System.Threading.Tasks
Imports TodoList.Models

Public Class IndexBuilder
    Implements IDisposable
    Private ReadOnly _fileExtractor As New FileExtractor()

    Private ReadOnly _index As New ConcurrentDictionary(Of String, FileItem)()
    Private _indexPath As String = Path.Combine(TaskInfo.appDirectory, "FileSearcher", "index.dat")

    Private supportedExtensions As String() = {".txt", ".log", ".md", ".csv", ".json", ".xml", ".html", ".htm",
                                          ".css", ".js", ".vb", ".cs", ".sql", ".docx", ".xlsx", ".pptx", ".pdf", ".doc", ".xls", ".ppt"}

    ' 可选：OCR 启用开关（也可以改为参数或配置）
    Private ReadOnly _enableOcr As Boolean = True

    ' 可选：OCR 限制（仅示例；具体阈值按需改）
    Private Const MaxOcrBytes As Long = 4L * 1024 * 1024 ' >4MB 的图片跳过 OCR

    ' 超时常量（按需调整）
    Private Const PerFileTimeoutMs As Integer = 120000 ' 每个文件最多 2 分钟

    ' ADD: 维护根目录集合（不区分大小写）
    Private ReadOnly _roots As New System.Collections.Concurrent.ConcurrentDictionary(Of String, Boolean)(StringComparer.OrdinalIgnoreCase)

    Private _cts As Threading.CancellationTokenSource

    Public Sub BeginNewOperation()
        _cts = New Threading.CancellationTokenSource()
    End Sub

    Public Sub RequestCancel()
        Try
            _cts?.Cancel()
        Catch
        End Try
    End Sub

    Private Function GetToken() As Threading.CancellationToken
        If _cts Is Nothing Then
            _cts = New Threading.CancellationTokenSource()
        End If
        Return _cts.Token
    End Function

    ' ADD: 规范化路径，去掉末尾分隔符
    Private Shared Function NormalizeRootPath(path As String) As String
        If String.IsNullOrWhiteSpace(path) Then Return path
        Dim full = System.IO.Path.GetFullPath(path)
        Return full.TrimEnd(System.IO.Path.DirectorySeparatorChar, System.IO.Path.AltDirectorySeparatorChar)
    End Function

    ' ADD: 判断文件是否位于某根目录下
    Private Shared Function IsUnderRoot(filePath As String, root As String) As Boolean
        Dim fp = NormalizeRootPath(filePath)
        Dim rt = NormalizeRootPath(root)
        If Not fp.StartsWith(rt, StringComparison.OrdinalIgnoreCase) Then Return False
        ' 确保是同一层级下（例如 C:\A 不是 C:\AB）
        If fp.Length = rt.Length Then Return True
        Dim ch = fp(rt.Length)
        Return ch = Path.DirectorySeparatorChar OrElse ch = Path.AltDirectorySeparatorChar
    End Function

    ' ADD: 对外暴露根目录集合只读视图
    Public Function GetIndexedRoots() As IReadOnlyCollection(Of String)
        Return _roots.Keys.ToList().AsReadOnly()
    End Function


    ' 软超时包装：成功则返回 True 和结果；超时返回 False
    Private Shared Function TryRunWithTimeout(Of T)(work As Func(Of T), timeoutMs As Integer, ByRef result As T) As Boolean
        Dim t1 = Task.Run(work)
        Dim ok = t1.Wait(timeoutMs)
        If ok Then
            result = t1.Result
            Return True
        Else
            Return False
        End If
    End Function

    Public Sub Dispose() Implements IDisposable.Dispose
        _fileExtractor.Dispose()
    End Sub

    ' 安全地递归枚举文件：遇到无法访问的目录时记录日志并跳过
    Private Iterator Function SafeEnumerateFiles(root As String) As IEnumerable(Of String)
        Dim stack As New Stack(Of String)()
        stack.Push(root)

        While stack.Count > 0
            Dim currentDir = stack.Pop()

            ' 先枚举子目录
            Dim subDirs As IEnumerable(Of String) = Enumerable.Empty(Of String)()
            Try
                subDirs = Directory.EnumerateDirectories(currentDir)
            Catch ex As UnauthorizedAccessException
                Logger.Log($"无法访问目录: {currentDir}，已跳过。原因: {ex.Message}")
                Continue While
            Catch ex As DirectoryNotFoundException
                Logger.Log($"目录不存在或已被删除: {currentDir}，已跳过。原因: {ex.Message}")
                Continue While
            Catch ex As PathTooLongException
                Logger.Log($"目录路径过长: {currentDir}，已跳过。原因: {ex.Message}")
                Continue While
            Catch ex As IOException
                Logger.Log($"读取目录出错: {currentDir}，已跳过。原因: {ex.Message}")
                Continue While
            End Try

            ' 压栈待处理
            For Each d In subDirs
                stack.Push(d)
            Next

            ' 再枚举当前目录下的文件
            Dim files As IEnumerable(Of String) = Enumerable.Empty(Of String)()
            Try
                files = Directory.EnumerateFiles(currentDir, "*.*", SearchOption.TopDirectoryOnly)
            Catch ex As UnauthorizedAccessException
                Logger.Log($"无法列出文件: {currentDir}，已跳过该目录文件。原因: {ex.Message}")
                files = Enumerable.Empty(Of String)()
            Catch ex As DirectoryNotFoundException
                Logger.Log($"目录不存在或已被删除: {currentDir}，已跳过该目录文件。原因: {ex.Message}")
                files = Enumerable.Empty(Of String)()
            Catch ex As PathTooLongException
                Logger.Log($"目录路径过长(文件枚举): {currentDir}，已跳过该目录文件。原因: {ex.Message}")
                files = Enumerable.Empty(Of String)()
            Catch ex As IOException
                Logger.Log($"读取目录文件出错: {currentDir}，已跳过该目录文件。原因: {ex.Message}")
                files = Enumerable.Empty(Of String)()
            End Try

            For Each f In files
                Yield f
            Next
        End While
    End Function


    ' 构建文件索引
    'Public Async Function BuildIndexAsync(rootDirectory As String, progress As IProgress(Of IndexingProgressInfo)) As Task(Of Integer)
    '    If Not Directory.Exists(rootDirectory) Then
    '        Throw New DirectoryNotFoundException($"目录不存在: {rootDirectory}")
    '    End If

    '    Dim lastReport As DateTime = DateTime.Now
    '    Dim reportCounter As Integer = 0
    '    Const reportEveryNFiles As Integer = 5
    '    Const reportEveryMs As Integer = 300

    '    ' 清空当前索引
    '    _index.Clear()

    '    ' 获取所有文件
    '    'Dim allFiles = Directory.GetFiles(rootDirectory, "*.*", SearchOption.AllDirectories)
    '    'Dim totalFiles As Integer = allFiles.Length
    '    Dim files As List(Of String) = SafeEnumerateFiles(rootDirectory).ToList()
    '    Dim totalFiles As Integer = files.Count
    '    Dim processedFiles As Integer = 0

    '    ' 报告初始进度
    '    progress.Report(New IndexingProgressInfo(0, 0, totalFiles, "准备索引..."))

    '    ' 获取所有文件（流式）
    '    Dim fileEnumerable = Directory.EnumerateFiles(rootDirectory, "*.*", SearchOption.AllDirectories)

    '    ' 为流式源创建分区器（不缓冲）
    '    'Dim partitioner1 = Partitioner.Create(fileEnumerable, EnumerablePartitionerOptions.NoBuffering)

    '    ' 并行处理：每个线程一个 FileExtractor
    '    Await Task.Run(
    '        Sub()
    '            Parallel.ForEach(Of String, FileExtractor)(
    '        files,'partitioner1,
    '        New ParallelOptions With {.MaxDegreeOfParallelism = Environment.ProcessorCount},
    '        Function() New FileExtractor(), ' localInit: 每线程实例
    '        Function(filePath, loopState, localExtractor)
    '            Try
    '                Dim fi As New FileInfo(filePath)
    '                Dim fileInfo = New FileItem() With {
    '                    .Path = filePath,
    '                    .Name = fi.Name,
    '                    .Size = fi.Length,
    '                    .Modified = fi.LastWriteTime,
    '                    .Extension = fi.Extension.ToLower()
    '                }

    '                ' 提取文本内容（文本文件 + 可选图片）
    '                Dim needContent As Boolean =
    '                    IsContentExtractable(fileInfo.Extension) OrElse
    '                    (_enableOcr AndAlso localExtractor._ocrProcessor.IsSupportedImageFile(filePath) AndAlso fi.Length <= MaxOcrBytes)

    '                If needContent Then
    '                    Dim content As String = Nothing
    '                    If TryRunWithTimeout(Function() localExtractor.ExtractTextContent(filePath), PerFileTimeoutMs, content) Then
    '                        fileInfo.Content = content
    '                        fileInfo.Keywords = localExtractor.ExtractKeywords(fileInfo.Content)
    '                    Else
    '                        Logger.Log($"提取超时，跳过: {filePath}")
    '                        ' 可选：打标记，方便后续增量补救
    '                        ' fileInfo.Content = ""
    '                        ' fileInfo.Keywords = New List(Of String)()
    '                    End If
    '                End If

    '                ' 添加到索引
    '                _index.TryAdd(filePath, fileInfo)

    '                ' 更新进度
    '                Dim currentProcessed = System.Threading.Interlocked.Increment(processedFiles)
    '                Dim percentComplete As Integer = CInt(Math.Floor((currentProcessed / CDbl(totalFiles)) * 100))

    '                ' 简单节流（轻度竞争可接受）
    '                reportCounter += 1
    '                If reportCounter >= reportEveryNFiles OrElse (DateTime.Now - lastReport).TotalMilliseconds > reportEveryMs Then
    '                    progress.Report(New IndexingProgressInfo(percentComplete, currentProcessed, totalFiles, Path.GetFileName(filePath)))
    '                    reportCounter = 0
    '                    lastReport = DateTime.Now
    '                End If
    '            Catch ex As Exception
    '                Logger.Log($"索引文件出错: {filePath}, 错误: {ex.Message}")
    '            End Try

    '            Return localExtractor ' 必须返回本地实例
    '        End Function,
    '        Sub(localExtractor)
    '            ' localFinally：释放每线程实例
    '            localExtractor.Dispose()
    '        End Sub)
    '        End Sub)

    '    ' 报告完成进度
    '    progress.Report(New IndexingProgressInfo(100, processedFiles, totalFiles, "完成"))

    '    ' 保存索引到磁盘
    '    SaveIndex()

    '    Return _index.Count
    'End Function

    ' CHANGE: 清空后，只对该根构建并保存
    Public Async Function BuildIndexAsync(rootDirectory As String, progress As IProgress(Of IndexingProgressInfo)) As Task(Of Integer)

        If Not Directory.Exists(rootDirectory) Then
            Throw New DirectoryNotFoundException($"目录不存在: {rootDirectory}")
        End If

        _index.Clear()
        _roots.Clear()

        Dim count = Await BuildIndexForRootAsync(rootDirectory, progress)
        _roots.TryAdd(NormalizeRootPath(rootDirectory), True)

        SaveIndex()
        Return count
    End Function

    ' ADD: 针对单一根目录的构建（不清空，累加）
    Public Async Function BuildIndexForRootAsync(rootDirectory As String, progress As IProgress(Of IndexingProgressInfo)) As Task(Of Integer)
        Dim token = GetToken()
        If Not Directory.Exists(rootDirectory) Then
            Throw New DirectoryNotFoundException($"目录不存在: {rootDirectory}")
        End If

        Dim lastReport As DateTime = DateTime.Now
        Dim reportCounter As Integer = 0
        Const reportEveryNFiles As Integer = 5
        Const reportEveryMs As Integer = 300

        Dim files As List(Of String) = SafeEnumerateFiles(rootDirectory).ToList()
        Dim totalFiles As Integer = files.Count
        Dim processedFiles As Integer = 0

        progress.Report(New IndexingProgressInfo(0, 0, totalFiles, "准备索引..."))

        Await Task.Run(Sub()
                           Parallel.ForEach(Of String, FileExtractor)(
            files,
            New ParallelOptions With {.MaxDegreeOfParallelism = Environment.ProcessorCount, .CancellationToken = token},
            Function() New FileExtractor(),
            Function(filePath, loopState, localExtractor)
                token.ThrowIfCancellationRequested()
                Try
                    Dim fi As New FileInfo(filePath)
                    Dim fileInfo = New FileItem() With {
                        .Path = filePath,
                        .Name = fi.Name,
                        .Size = fi.Length,
                        .Modified = fi.LastWriteTime,
                        .Extension = fi.Extension.ToLower()
                    }

                    Dim needContent As Boolean =
                        IsContentExtractable(fileInfo.Extension) OrElse
                        (_enableOcr AndAlso localExtractor._ocrProcessor.IsSupportedImageFile(filePath) AndAlso fi.Length <= MaxOcrBytes)

                    If needContent Then
                        Dim content As String = Nothing
                        If TryRunWithTimeout(Function() localExtractor.ExtractTextContent(filePath), PerFileTimeoutMs, content) Then
                            fileInfo.Content = content
                            fileInfo.Keywords = localExtractor.ExtractKeywords(fileInfo.Content)
                        Else
                            Logger.Log($"提取超时，跳过: {filePath}")
                        End If
                    End If

                    _index(filePath) = fileInfo

                    Dim currentProcessed = Interlocked.Increment(processedFiles)
                    Dim percentComplete As Integer = CInt(Math.Floor((currentProcessed / CDbl(totalFiles)) * 100))

                    reportCounter += 1
                    If reportCounter >= reportEveryNFiles OrElse (DateTime.Now - lastReport).TotalMilliseconds > reportEveryMs Then
                        progress.Report(New IndexingProgressInfo(percentComplete, currentProcessed, totalFiles, Path.GetFileName(filePath)))
                        reportCounter = 0
                        lastReport = DateTime.Now
                    End If
                Catch ex As Exception
                    Logger.Log($"索引文件出错: {filePath}, 错误: {ex.Message}")
                End Try

                Return localExtractor
            End Function,
            Sub(localExtractor) localExtractor.Dispose()
        )
                       End Sub, token)

        progress.Report(New IndexingProgressInfo(100, processedFiles, totalFiles, "完成"))

        ' 不在这里保存，交由上层调用时决定（但可按需保存）
        Return processedFiles
    End Function


    Private Function ExtractWithWorker(filePath As String, timeoutMs As Integer) As String
        Dim psi As New ProcessStartInfo() With {
        .FileName = Path.Combine(AppDomain.CurrentDomain.BaseDirectory, "FileExtractor.Worker.exe"),
        .Arguments = """" & filePath & """",
        .UseShellExecute = False,
        .CreateNoWindow = True,
        .RedirectStandardOutput = True,
        .RedirectStandardError = True
    }

        Using p As New Process()
            p.StartInfo = psi
            p.Start()
            If p.WaitForExit(timeoutMs) Then
                Dim output As String = p.StandardOutput.ReadToEnd()
                ' 可选：读取错误输出便于诊断
                ' Dim err As String = p.StandardError.ReadToEnd()
                Return output
            Else
                Try
                    p.Kill(entireProcessTree:=True)
                Catch
                End Try
                Logger.Log($"工作进程超时，已终止: {filePath}")
                Return String.Empty
            End If
        End Using
    End Function

    ' 确定是否支持内容提取
    Private Function IsContentExtractable(extension As String) As Boolean
        Return supportedExtensions.Contains(extension)
    End Function

    ' 增量更新索引
    'Public Async Function UpdateIndexAsync(rootDirectory As String, progress As IProgress(Of IndexingProgressInfo)) As Task(Of Integer)
    '    If Not Directory.Exists(rootDirectory) Then
    '        Throw New DirectoryNotFoundException($"目录不存在: {rootDirectory}")
    '    End If

    '    Dim lastReport As DateTime = DateTime.Now

    '    Dim reportCounter As Integer = 0
    '    Const reportEveryNFiles As Integer = 5
    '    Const reportEveryMs As Integer = 300

    '    ' 加载现有索引
    '    LoadIndex()

    '    ' 获取所有文件
    '    Dim allFiles = Directory.GetFiles(rootDirectory, "*.*", SearchOption.AllDirectories).ToHashSet()
    '    Dim filesToProcess As New List(Of String)()
    '    Dim filesToRemove As New List(Of String)()

    '    ' 找出需要从索引中删除的文件 (不再存在的文件)
    '    For Each indexedPath In _index.Keys
    '        If Not allFiles.Contains(indexedPath) Then
    '            filesToRemove.Add(indexedPath)
    '        End If
    '    Next

    '    ' 找出需要添加或更新的文件
    '    For Each filePath In allFiles
    '        ' 如果文件不在索引中，或者文件已被修改，则添加到处理列表
    '        If Not _index.ContainsKey(filePath) OrElse
    '           File.GetLastWriteTime(filePath) > _index(filePath).Modified Then
    '            filesToProcess.Add(filePath)
    '        End If
    '    Next

    '    ' 从索引中删除不存在的文件
    '    For Each path In filesToRemove
    '        Dim removed As FileItem = Nothing
    '        _index.TryRemove(path, removed)
    '    Next

    '    ' 处理新文件或已修改的文件
    '    Dim totalToProcess As Integer = filesToProcess.Count
    '    Dim processed As Integer = 0

    '    ' 报告初始进度
    '    progress.Report(New IndexingProgressInfo(0, 0, totalToProcess, "准备更新..."))

    '    Await Task.Run(Sub()
    '                       Dim partitioner1 = Partitioner.Create(filesToProcess, EnumerablePartitionerOptions.NoBuffering)
    '                       Parallel.ForEach(Of String, FileExtractor)(
    '                        partitioner1,
    '                        New ParallelOptions With {.MaxDegreeOfParallelism = Environment.ProcessorCount},
    '                        Function() New FileExtractor(),
    '                        Function(filePath, loopState, localExtractor)
    '                            Try
    '                                Dim fi As New FileInfo(filePath)
    '                                Dim fileInfo = New FileItem() With {
    '                                    .Path = filePath,
    '                                    .Name = fi.Name,
    '                                    .Size = fi.Length,
    '                                    .Modified = fi.LastWriteTime,
    '                                    .Extension = fi.Extension.ToLower()
    '                                }

    '                                Dim needContent As Boolean = IsContentExtractable(fileInfo.Extension) OrElse
    '                                    (_enableOcr AndAlso localExtractor._ocrProcessor.IsSupportedImageFile(filePath) AndAlso fi.Length <= MaxOcrBytes)

    '                                If needContent Then
    '                                    Dim content As String = Nothing
    '                                    If TryRunWithTimeout(Function() localExtractor.ExtractTextContent(filePath), PerFileTimeoutMs, content) Then
    '                                        fileInfo.Content = content
    '                                        fileInfo.Keywords = localExtractor.ExtractKeywords(fileInfo.Content)
    '                                    Else
    '                                        Logger.Log($"提取超时，跳过: {filePath}")
    '                                        ' 可选：打标记，方便后续增量补救
    '                                        ' fileInfo.Content = ""
    '                                        ' fileInfo.Keywords = New List(Of String)()
    '                                    End If
    '                                End If

    '                                _index(filePath) = fileInfo

    '                                Dim currentProcessed = System.Threading.Interlocked.Increment(processed)
    '                                Dim percentComplete As Integer = CInt(Math.Floor((currentProcessed / CDbl(totalToProcess)) * 100))

    '                                reportCounter += 1
    '                                If reportCounter >= reportEveryNFiles OrElse (DateTime.Now - lastReport).TotalMilliseconds > reportEveryMs Then
    '                                    progress.Report(New IndexingProgressInfo(percentComplete, currentProcessed, totalToProcess, Path.GetFileName(filePath)))
    '                                    reportCounter = 0
    '                                    lastReport = DateTime.Now
    '                                End If

    '                            Catch ex As Exception
    '                                Logger.Log($"更新索引出错: {filePath}, 错误: {ex.Message}")
    '                            End Try

    '                            Return localExtractor
    '                        End Function,
    '        Sub(localExtractor) localExtractor.Dispose()
    '    )
    '                   End Sub)


    '    ' 报告完成进度
    '    progress.Report(New IndexingProgressInfo(100, processed, totalToProcess, "完成"))

    '    ' 保存更新后的索引
    '    SaveIndex()

    '    Return filesToProcess.Count
    'End Function

    ' CHANGE: 仅更新单一根目录，不影响其它根
    Public Async Function UpdateIndexForRootAsync(rootDirectory As String, progress As IProgress(Of IndexingProgressInfo)) As Task(Of Integer)
        Dim token = GetToken()

        If Not Directory.Exists(rootDirectory) Then
            Throw New DirectoryNotFoundException($"目录不存在: {rootDirectory}")
        End If

        Dim rootNorm = NormalizeRootPath(rootDirectory)

        ' 注意：不在这里 LoadIndex()，由调用方在程序启动时统一加载
        Dim lastReport As DateTime = DateTime.Now
        Dim reportCounter As Integer = 0
        Const reportEveryNFiles As Integer = 5
        Const reportEveryMs As Integer = 300

        ' 获取该根目录下的文件全集
        Dim allFiles = Directory.GetFiles(rootNorm, "*.*", SearchOption.AllDirectories).ToHashSet(StringComparer.OrdinalIgnoreCase)

        ' 找出需要从索引中删除的文件（仅限该根目录下）
        Dim filesToRemove As New List(Of String)()
        For Each indexedPath In _index.Keys
            If IsUnderRoot(indexedPath, rootNorm) AndAlso Not allFiles.Contains(indexedPath) Then
                filesToRemove.Add(indexedPath)
            End If
        Next

        ' 找出需要添加或更新的文件（仅限该根目录下）
        Dim filesToProcess As New List(Of String)()
        For Each filePath In allFiles
            If Not _index.ContainsKey(filePath) OrElse
           File.GetLastWriteTime(filePath) > _index(filePath).Modified Then
                filesToProcess.Add(filePath)
            End If
        Next

        ' 执行删除
        For Each path In filesToRemove
            Dim removed As FileItem = Nothing
            _index.TryRemove(path, removed)
        Next

        ' 处理新文件或已修改的文件
        Dim totalToProcess As Integer = filesToProcess.Count
        Dim processed As Integer = 0

        progress.Report(New IndexingProgressInfo(0, 0, totalToProcess, "准备更新..."))

        Await Task.Run(Sub()
                           Dim partitioner1 = Partitioner.Create(filesToProcess, EnumerablePartitionerOptions.NoBuffering)
                           Parallel.ForEach(Of String, FileExtractor)(
            partitioner1,
            New ParallelOptions With {.MaxDegreeOfParallelism = Environment.ProcessorCount, .CancellationToken = token},
            Function() New FileExtractor(),
            Function(filePath, loopState, localExtractor)
                token.ThrowIfCancellationRequested()
                Try
                    Dim fi As New FileInfo(filePath)
                    Dim fileInfo = New FileItem() With {
                        .Path = filePath,
                        .Name = fi.Name,
                        .Size = fi.Length,
                        .Modified = fi.LastWriteTime,
                        .Extension = fi.Extension.ToLower()
                    }

                    Dim needContent As Boolean =
                        IsContentExtractable(fileInfo.Extension) OrElse
                        (_enableOcr AndAlso localExtractor._ocrProcessor.IsSupportedImageFile(filePath) AndAlso fi.Length <= MaxOcrBytes)

                    If needContent Then
                        Dim content As String = Nothing
                        If TryRunWithTimeout(Function() localExtractor.ExtractTextContent(filePath), PerFileTimeoutMs, content) Then
                            fileInfo.Content = content
                            fileInfo.Keywords = localExtractor.ExtractKeywords(fileInfo.Content)
                        Else
                            Logger.Log($"提取超时，跳过: {filePath}")
                        End If
                    End If

                    _index(filePath) = fileInfo

                    Dim currentProcessed = Interlocked.Increment(processed)
                    Dim percentComplete As Integer = CInt(Math.Floor((currentProcessed / CDbl(totalToProcess)) * 100))

                    reportCounter += 1
                    If reportCounter >= reportEveryNFiles OrElse (DateTime.Now - lastReport).TotalMilliseconds > reportEveryMs Then
                        progress.Report(New IndexingProgressInfo(percentComplete, currentProcessed, totalToProcess, Path.GetFileName(filePath)))
                        reportCounter = 0
                        lastReport = DateTime.Now
                    End If

                Catch ex As Exception
                    Logger.Log($"更新索引出错: {filePath}, 错误: {ex.Message}")
                End Try

                Return localExtractor
            End Function,
            Sub(localExtractor) localExtractor.Dispose()
        )
                       End Sub, token)

        progress.Report(New IndexingProgressInfo(100, processed, totalToProcess, "完成"))

        SaveIndex()
        Return filesToProcess.Count
    End Function

    ' ADD: “添加索引”（若已存在则更新）
    Public Async Function AddOrUpdateRootAsync(rootDirectory As String, progress As IProgress(Of IndexingProgressInfo)) As Task(Of Integer)
        If Not Directory.Exists(rootDirectory) Then
            Throw New DirectoryNotFoundException($"目录不存在: {rootDirectory}")
        End If

        Dim rootNorm = NormalizeRootPath(rootDirectory)
        Dim processed As Integer

        If _roots.ContainsKey(rootNorm) Then
            processed = Await UpdateIndexForRootAsync(rootNorm, progress)
        Else
            processed = Await BuildIndexForRootAsync(rootNorm, progress)
            _roots.TryAdd(rootNorm, True)
            SaveIndex()
        End If

        Return processed
    End Function

    ' ADD: 新的持久化对象
    Public Class PersistedIndexData
        Public Property Roots As List(Of String)
        Public Property Items As List(Of KeyValuePairData)
    End Class


    Private Function ExtractTextContent(filePath As String) As String
        Return _fileExtractor.ExtractTextContent(filePath)
    End Function

    Private Function ExtractKeywords(content As String) As List(Of String)
        Return _fileExtractor.ExtractKeywords(content)
    End Function

    ' 保存索引到磁盘
    'Public Sub SaveIndex()
    '    Try
    '        ' 确保目录存在
    '        Dim directory1 = Path.GetDirectoryName(_indexPath)
    '        If Not Directory.Exists(directory1) Then
    '            Directory.CreateDirectory(directory1)
    '        End If

    '        ' 将索引数据转换为可序列化的格式
    '        Dim serializableData = _index.Select(Function(kv) New KeyValuePairData With {
    '            .Key = kv.Key,
    '            .Value = kv.Value
    '        }).ToList()

    '        ' JSON 序列化
    '        Dim options = New JsonSerializerOptions With {
    '            .WriteIndented = False ' 关闭缩进，减少写盘与 CPU 时间
    '        }
    '        Dim jsonString = JsonSerializer.Serialize(serializableData, options)
    '        File.WriteAllText(_indexPath, jsonString)

    '    Catch ex As Exception
    '        Logger.Log($"保存索引出错: {ex.Message}")
    '    End Try
    'End Sub

    ' CHANGE: 保存索引（包含根目录）
    Public Sub SaveIndex()
        Try
            Dim directory1 = Path.GetDirectoryName(_indexPath)
            If Not Directory.Exists(directory1) Then
                Directory.CreateDirectory(directory1)
            End If

            Dim data As New PersistedIndexData With {
            .Roots = _roots.Keys.ToList(),
            .Items = _index.Select(Function(kv) New KeyValuePairData With {.Key = kv.Key, .Value = kv.Value}).ToList()
        }

            Dim options = New JsonSerializerOptions With {.WriteIndented = False}
            Dim jsonString = JsonSerializer.Serialize(data, options)
            File.WriteAllText(_indexPath, jsonString)

        Catch ex As Exception
            Logger.Log($"保存索引出错: {ex.Message}")
        End Try
    End Sub


    '' 从磁盘加载索引
    'Public Sub LoadIndex()
    '    Try
    '        If File.Exists(_indexPath) Then
    '            ' JSON 反序列化
    '            Dim jsonString = File.ReadAllText(_indexPath)
    '            Dim serializableData = JsonSerializer.Deserialize(Of List(Of KeyValuePairData))(jsonString)

    '            _index.Clear()
    '            For Each item In serializableData
    '                _index.TryAdd(item.Key, item.Value)
    '            Next
    '        End If
    '    Catch ex As Exception
    '        Logger.Log($"加载索引出错: {ex.Message}")
    '        _index.Clear() ' 如果加载失败，清空索引
    '    End Try
    'End Sub

    ' CHANGE: 加载索引（兼容旧格式）
    Public Sub LoadIndex()
        Try
            If File.Exists(_indexPath) Then
                Dim jsonString = File.ReadAllText(_indexPath)

                Dim data As PersistedIndexData = Nothing
                Dim ok As Boolean = False
                Try
                    data = JsonSerializer.Deserialize(Of PersistedIndexData)(jsonString)
                    ok = (data IsNot Nothing AndAlso data.Items IsNot Nothing)
                Catch
                    ok = False
                End Try

                _index.Clear()
                _roots.Clear()

                If ok Then
                    For Each item In data.Items
                        _index.TryAdd(item.Key, item.Value)
                    Next
                    If data.Roots IsNot Nothing Then
                        For Each r In data.Roots
                            _roots.TryAdd(NormalizeRootPath(r), True)
                        Next
                    End If
                Else
                    ' 兼容旧格式：旧文件只有 Items 数组
                    Dim serializableData = JsonSerializer.Deserialize(Of List(Of KeyValuePairData))(jsonString)
                    If serializableData IsNot Nothing Then
                        For Each item In serializableData
                            _index.TryAdd(item.Key, item.Value)
                        Next
                    End If
                    ' Roots 不存在（留空）；后续第一次“添加索引”会补齐
                End If
            End If
        Catch ex As Exception
            Logger.Log($"加载索引出错: {ex.Message}")
            _index.Clear()
            _roots.Clear()
        End Try
    End Sub

    ' 用于 JSON 序列化的辅助类
    Public Class KeyValuePairData
        Public Property Key As String
        Public Property Value As FileItem
    End Class
    ' 获取索引统计信息
    Public Function GetIndexStats() As (TotalFiles As Integer, TotalSize As Long, IndexedDate As DateTime?)
        Dim indexedDate As DateTime? = Nothing

        If File.Exists(_indexPath) Then
            indexedDate = File.GetLastWriteTime(_indexPath)
        End If

        Dim totalSize As Long = _index.Values.Sum(Function(f) f.Size)

        Return (_index.Count, totalSize, indexedDate)
    End Function

    ' 获取当前索引
    Public Function GetIndex() As IReadOnlyDictionary(Of String, FileItem)
        Return _index
    End Function
End Class
